import pandas as pd

from app_config import get_engine_ts, get_pro
import numpy as np
import time

# 一年前
start_date_str = '20221101'
# 一季度前
end_date_str = '20231031'
# 当前时间
kc_start_date_str = "20230801"
# 市盈率计算日
ttm_date_str = "20231031"

counter = {'count': 0}


def get_profit(st_code):
    counter['count'] += 1
    print(counter['count'])
    if counter['count'] % 160 == 0:
        print(f"Sleeping for 60 seconds after 150 requests...")
        time.sleep(62)

    df_ = get_pro().query('fina_indicator', ts_code=st_code,
                          fields="ts_code,end_date,q_opincome,q_investincome,q_dtprofit,q_eps,q_netprofit_margin,"
                                 "q_gsprofit_margin,q_exp_to_sales,q_profit_to_gr,q_saleexp_to_gr,q_adminexp_to_gr,"
                                 "q_finaexp_to_gr,q_impair_to_gr_ttm,q_gc_to_gr,q_op_to_gr,q_roe,q_dt_roe,q_npta,"
                                 "q_opincome_to_ebt,q_investincome_to_ebt,q_dtprofit_to_profit,q_salescash_to_or,"
                                 "q_ocf_to_sales,q_ocf_to_or", start_date=start_date_str,
                          end_date=end_date_str)

    df_['q_profit'] = df_['q_dtprofit'] / df_['q_dtprofit_to_profit'] * 100
    print("财务数据 " + st_code + "  :  " + str(len(df_)))
    if len(df_) > 4:
        print("err =================" + str(len(df_)))
    df_dedup = df_.drop_duplicates(subset='end_date', keep='first')
    total_q_dtprofit = df_dedup['q_dtprofit'].sum()
    total_q_profit = df_dedup['q_profit'].sum()
    return total_q_dtprofit / 100000000, total_q_profit / 100000000


def get_profit1(st_code):
    return 100.0, 200.0


def calculate_volatility(df):
    df = df.sort_values(by='trade_date')

    df_one_year = df.copy()  # 使用.copy()显式创建副本

    # 计算每日对数收益率
    df_one_year['log_return'] = np.log(df_one_year['close'] / df_one_year['close'].shift(1))

    # 删除缺失值
    df_one_year = df_one_year.dropna(subset=['log_return'])

    # 计算波动率（年化）
    volatility = df_one_year['log_return'].std() * np.sqrt(len(df_one_year))  # 252是交易日的近似值

    return volatility


if __name__ == '__main__':
    from _stock_basic import get_data

    engine = get_engine_ts()

    # # 获取今天的日期
    # today = datetime.now()
    """     1 样本空间    """
    ''' 1.1 上市时间超过 12 个月； '''
    stock_basic = get_data()

    print("全部股票:")
    print(len(stock_basic))

    ''' 1.2 被实施退市风险警示除外； '''
    # 删除 name 列包含 'ST' 字符串的行
    st_filtered = stock_basic[~stock_basic['name'].str.contains('ST')]
    print("排除:")
    print(len(st_filtered))

    # 根据条件过滤数据
    filtered_df = st_filtered[(
                                      (st_filtered['market'].isin(['科创板', '创业板']))
                                      &
                                      (st_filtered['list_date'] <= start_date_str)
                              ) | (
                                      (st_filtered['market'] == '主板')
                                      &
                                      (st_filtered['list_date'] <= kc_start_date_str)
                              )]
    print("样本数量")
    print(len(filtered_df))

    """     2、选样方法    """
    '''2.1对样本空间内的证券按照过去一年的日均成交金额由高到低排名，剔除排名后 10%的证券作为待选样本；'''
    # 转换 ts_code 列为元组，并生成 SQL 中 IN 子句所需的格式
    ts_codes = filtered_df['ts_code'].tolist()

    # 构建查询语句
    query = f"""
    SELECT * FROM `daily_basic`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
    AND trade_date >= '{start_date_str}'
    AND trade_date <= '{end_date_str}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)

    result_df['amount'] = result_df['circ_mv'] * result_df['turnover_rate']

    '''对待选样本按照过去一年的日均总市值由高到低排名，选取排名前 50的证券作为指数样本。'''

    # 以 ts_code 列分组，并计算 total_mv 列的平均值
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', 'avg_amount']
    merge = pd.merge(filtered_df, grouped_df, on='ts_code', how='left')

    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', 'avg_total_mv']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    # 按 avg_amount 列降序排序
    df_sorted = pd_merge.sort_values(by='avg_amount', ascending=False)
    # 获取前 50% 的数据
    num_rows = len(df_sorted)
    top_50_percent = df_sorted.head(int(num_rows * 0.5))

    top_50_percent = top_50_percent.sort_values('avg_total_mv', ascending=False)
    top_50_percent = top_50_percent.head(300)
    # top_50_percent.to_excel("top_300.xlsx")

    '''计算年化波动率 好像没什么用'''
    # 测试函数（需根据实际情景选择相应的DataFrame）
    # 假设我们已经加载了数据在df中，并且它是按ts_code分组的
    volatility_dict = {}
    # 遍历计算波动率
    for stock_code, group_df in result_df.groupby('ts_code'):
        volatility_dict[stock_code] = calculate_volatility(group_df)
    # 将字典转换为 DataFrame
    volatility_df = pd.DataFrame(list(volatility_dict.items()), columns=['ts_code', 'volatility'])

    d = pd.merge(top_50_percent, volatility_df, on='ts_code', how='left')

    '''计算TTM'''
    # 使用 apply 方法添加两列 total_q_dtprofit 和 total_q_profit
    # todo
    d[['total_q_dtprofit', 'total_q_profit']] = d['ts_code'].apply(lambda _ts_code: pd.Series(get_profit(_ts_code)))

    # 构建查询语句
    query_date = f"""
      SELECT ts_code, total_mv FROM `daily_basic`
      WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
      AND trade_date = '{ttm_date_str}'
      """

    # 执行查询并将结果转换为DataFrame
    result_df_date = pd.read_sql_query(query_date, engine)
    result_df_date['total_mv'] = result_df_date['total_mv'] / 10000

    data_frame_d = pd.merge(d, result_df_date, on='ts_code', how='left')

    # 确认列名无冲突
    print(data_frame_d.columns)
    data_frame_d['ttm'] = data_frame_d['total_mv'] / data_frame_d['total_q_profit']

    data_frame_d['dt_ttm'] = data_frame_d['total_mv'] / data_frame_d['total_q_dtprofit']

    data_frame_d.to_excel('top_300_2023.xlsx')



